FEAT Add BeaverTails dataset loader by romanlutz · Pull Request #1424 · Azure/PyRIT

romanlutz · 2026-03-01T14:28:01Z

Add remote dataset loader for BeaverTails (PKU-Alignment/BeaverTails), containing 330k+ QA pairs annotated across 14 harm categories for safety alignment research. Filters to unsafe entries by default.

Copilot

Pull request overview

Adds a new remote seed dataset loader for the BeaverTails HuggingFace dataset, making it discoverable via SeedDatasetProvider and documenting its availability.

Changes:

Introduces _BeaverTailsDataset remote loader with optional unsafe_only filtering (default: unsafe only).
Registers the loader in the remote datasets module and adds unit tests for filtering behavior.
Updates the “Loading Built-in Datasets” notebook output to include the new dataset name.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
`pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py`	New HuggingFace-backed loader that converts BeaverTails rows into `SeedPrompt`s (unsafe-only by default).
`pyrit/datasets/seed_datasets/remote/__init__.py`	Imports/exports the new loader so it’s auto-registered/discoverable.
`tests/unit/datasets/test_beaver_tails_dataset.py`	Adds unit tests covering unsafe-only vs all-entries behavior and dataset naming.
`doc/code/datasets/1_loading_datasets.ipynb`	Notebook updated to reflect the new dataset in the available list (but now includes executed outputs/metadata).

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py

doc/code/datasets/1_loading_datasets.ipynb

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.

doc/code/datasets/1_loading_datasets.ipynb

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py

doc/code/datasets/1_loading_datasets.ipynb

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Add remote dataset loader for BeaverTails (PKU-Alignment/BeaverTails), containing 330k+ QA pairs annotated across 14 harm categories for safety alignment research. Filters to unsafe entries by default. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The HF dataset identifier is now a class constant HF_DATASET_NAME instead of a constructor parameter, consistent with other loaders. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

For a 330k-row dataset, this avoids hundreds of thousands of redundant string/list allocations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

Copilot · 2026-03-02T21:56:51Z

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py

+        description = (
+            "BeaverTails is a collection of 330k+ human-LLM QA pairs annotated across 14 harm "
+            "categories, designed for safety alignment research. Introduced in 'BeaverTails: "
+            "Towards Improved Safety Alignment of LLM via a Human-Preference Dataset' (2023)."
+        )
+
+        source_url = f"https://huggingface.co/datasets/{self.HF_DATASET_NAME}"
+        groups = ["Institute for Artificial Intelligence", "CFCS, School of Computer Science"]
+
+        seed_prompts = []
+        for item in data:
+            if self.unsafe_only and item["is_safe"]:
+                continue
+
+            harm_categories = [k for k, v in item["category"].items() if v]
+
+            seed_prompts.append(
+                SeedPrompt(
+                    value=f"{{% raw %}}{item['prompt']}{{% endraw %}}",
+                    data_type="text",
+                    dataset_name=self.dataset_name,
+                    harm_categories=harm_categories,
+                    description=description,
+                    source=source_url,
+                    authors=authors,
+                    groups=groups,


The description/docstring emphasizes that BeaverTails contains QA pairs, but the loader currently only emits SeedPrompt values from item['prompt'] and ignores the associated response. To avoid misleading consumers, either (a) explicitly document that only the prompt column is extracted (similar to other dataset loaders), or (b) include the response in SeedPrompt.metadata (or a paired seed type if supported) so the QA relationship isn’t lost.

Copilot · 2026-03-02T21:56:52Z

doc/code/datasets/1_loading_datasets.ipynb

    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
-      "C:\\Users\\romanlutz\\AppData\\Local\\Temp\\ipykernel_50620\\4021500943.py:10: DeprecationWarning: is_objective parameter is deprecated since 0.13.0. Use seed_type='objective' instead.\n",
+      "C:\\Users\\romanlutz\\AppData\\Local\\Temp\\ipykernel_50556\\4021500943.py:10: DeprecationWarning: is_objective parameter is deprecated since 0.13.0. Use seed_type='objective' instead.\n",
      "  memory.get_seeds(harm_categories=[\"illegal\"], is_objective=True)\n"
     ]


This notebook diff still includes captured runtime output with user/machine-specific absolute paths (e.g., C:\\Users\\...\\AppData\\Local\\Temp\\ipykernel_...). Please clear cell outputs (and any execution metadata) before committing so docs remain deterministic and don’t leak local environment details.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI review requested due to automatic review settings March 1, 2026 14:28

Copilot started reviewing on behalf of romanlutz March 1, 2026 14:28 View session

romanlutz force-pushed the romanlutz/add-beaver-tails-dataset branch from 7b635d9 to b652d70 Compare March 1, 2026 14:28

Copilot AI reviewed Mar 1, 2026

View reviewed changes

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py Outdated Show resolved Hide resolved

doc/code/datasets/1_loading_datasets.ipynb Outdated Show resolved Hide resolved

doc/code/datasets/1_loading_datasets.ipynb Show resolved Hide resolved

romanlutz force-pushed the romanlutz/add-beaver-tails-dataset branch 2 times, most recently from 9741ae3 to 1fd2ef7 Compare March 2, 2026 13:02

Copilot AI review requested due to automatic review settings March 2, 2026 13:02

Copilot started reviewing on behalf of romanlutz March 2, 2026 13:03 View session

Copilot AI reviewed Mar 2, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings March 2, 2026 13:56

Copilot started reviewing on behalf of romanlutz March 2, 2026 13:57 View session

Copilot AI reviewed Mar 2, 2026

View reviewed changes

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py Outdated Show resolved Hide resolved

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py Show resolved Hide resolved

doc/code/datasets/1_loading_datasets.ipynb Show resolved Hide resolved

Copilot AI review requested due to automatic review settings March 2, 2026 15:07

Copilot started reviewing on behalf of romanlutz March 2, 2026 15:07 View session

Copilot AI reviewed Mar 2, 2026

View reviewed changes

romanlutz and others added 7 commits March 2, 2026 13:48

Remove dataset_name from constructor, hardcode as class constant

3a71604

The HF dataset identifier is now a class constant HF_DATASET_NAME instead of a constructor parameter, consistent with other loaders. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Use AsyncMock for _fetch_from_huggingface in tests

4f4fe8d

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Precompute source_url and groups outside the loop

9fd9044

For a 330k-row dataset, this avoids hundreds of thousands of redundant string/list allocations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Wrap prompt values in raw/endraw to preserve Jinja2 syntax

f98493c

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add license notice and content warning to docstring

e8d1379

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

fix: update notebook output for rebased datasets

a91052f

Copilot AI review requested due to automatic review settings March 2, 2026 21:50

romanlutz force-pushed the romanlutz/add-beaver-tails-dataset branch from 8a9dccb to a91052f Compare March 2, 2026 21:50

Copilot started reviewing on behalf of romanlutz March 2, 2026 21:51 View session

Copilot AI reviewed Mar 2, 2026

View reviewed changes

merge main, add E402/E501 to doc per-file-ignores

29dec57

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FEAT Add BeaverTails dataset loader#1424

FEAT Add BeaverTails dataset loader#1424
romanlutz wants to merge 8 commits intoAzure:mainfrom
romanlutz:romanlutz/add-beaver-tails-dataset

romanlutz commented Mar 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 2, 2026

Uh oh!

Copilot AI Mar 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

romanlutz commented Mar 1, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants